{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Test Import in Colab\n",
    "\n",
    "\"Run all\" to test that all the Flax imports work in head.\n",
    "\n",
    "Change runtime type as needed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Colab runtimes are pre-built with JAX/Flax:\n",
    "!pip freeze | egrep 'jax|flax'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "skip-execution"
    ]
   },
   "outputs": [],
   "source": [
    "# Install from head\n",
    "!pip install git+https://github.com/google/flax.git"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check versions after installing Flax from Github:\n",
    "!pip freeze | egrep 'jax|flax'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Verify we can import everything.\n",
    "import flax\n",
    "from flax.training import (checkpoints, dynamic_scale, early_stopping, lr_schedule,\n",
    "                           orbax_utils, prefetch_iterator, train_state, common_utils)\n",
    "from flax.metrics import tensorboard"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
